import numpy as np
from tensorflow.keras import datasets, layers, models
import matplotlib.pyplot as plt

# 加载Fashion MNIST数据集
(train_images, train_labels), (test_images, test_labels) = datasets.fashion_mnist.load_data()

# 类别名称
name = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'trou', 'Bag', 'Ankle boot']

# 图像归一化
train_images = train_images / 255.0
test_images = test_images / 255.0

# 为CNN添加颜色通道
train_images = train_images.reshape((train_images.shape[0], 28, 28, 1))
test_images = test_images.reshape((test_images.shape[0], 28, 28, 1))

# 定义序列模型model
model = (models.Sequential
([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),  # 第一层卷积层，Relu激活函数
    layers.MaxPooling2D((2, 2)),  # 最大池化层，降低图的空间尺寸
    layers.Conv2D(64, (3, 3), activation='relu'),  # 第二层卷积层
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),  # 将多维数据展平为一维，为全连接层做准备
    layers.Dense(64, activation='relu'),  # 全连接层
    layers.Dense(10, activation='softmax')  # 输出层，softmax激活函数输出分类概率分布
]))

# 编译模型
model.compile(optimizer='adam',  # 优化器
              loss='sparse_categorical_crossentropy',  # 损失函数为稀疏分类交叉熵
              metrics=['accuracy'])  # 评估指标为准确率

# 训练模型
history = model.fit(train_images, train_labels,
                    epochs=10,  # 十轮
                    validation_data=(test_images, test_labels))

# 测试
test_loss, test_acc = model.evaluate(test_images, test_labels)
print(f"Test accuracy: {test_acc}")  # f-string 格式化输出

# 绘制训练和测试的准确率曲线
plt.plot(history.history['accuracy'], label='accuracy')  # 训练集准确率曲线
plt.plot(history.history['val_accuracy'], label = 'val_accuracy')  # 测试集
plt.xlabel('Epoch')  # x轴标签
plt.ylabel('Accuracy')
plt.ylim([0, 1])  # y轴的取值范围为0到1
plt.legend(loc='lower right')  # 设置图例显示在右下角
plt.show()

pre = model.predict(test_images)

# 显示图像及其预测结果
def plot_image(i, prearray, trlabel, img):  # 获取第i个样本的预测数组、真实标签和图像
    prearray, trlabel, img = prearray[i], trlabel[i], img[i]
    plt.grid(False)  # 关闭网格
    plt.xticks([])   # 移除x轴刻度
    plt.yticks([])
    plt.imshow(img[...,0], cmap=plt.cm.binary)  # 显示图像，使用二值化颜色映射
    prelabel = np.argmax(prearray)
    if prelabel == trlabel:
        color = 'blue'
    else:
        color = 'red'

    plt.xlabel("{} {:2.0f}% ({})".format(name[prelabel],
                                         100*np.max(prearray),
                                         name[trlabel]),
                                         color=color)   # f-string

# 显示各个类别的预测概率
def plot_value_array(i, prearray, trlabel):
    prearray, trlabel = prearray[i], trlabel[i]
    plt.grid(False)
    plt.xticks(range(10))  # 设置x轴刻度为0-9
    plt.yticks([])
    thisplot = plt.bar(range(10), prearray, color="#777777")  # 绘制条形图，显示每个类别的预测概率，颜色设置为灰色
    plt.ylim([0, 1])
    prelabel = np.argmax(prearray)

    thisplot[prelabel].set_color('red')  # 将预测标签对应的条形图颜色设为红色
    thisplot[trlabel].set_color('blue')  # 将真实标签对应的条形图颜色设为蓝色

# 绘制测试图像及其预测结果
rows = 5  # 图像行数
cols = 3  # 列数
num = rows * cols  # 总数
plt.figure(figsize=(2 * 2 * cols, 2 * rows))  # 创建一个新的图形，设置图形大小
for i in range(num):
    plt.subplot(rows, 2 * cols, 2 * i + 1)
    plot_image(i, pre, test_labels, test_images)
    plt.subplot(rows, 2 * cols, 2 * i + 2)
    plot_value_array(i, pre, test_labels)
plt.tight_layout()
plt.show()

model.save('model.keras')